// 方法一：遍历 + 计数
export function countSubarrays(nums: number[], k: number): number {
    const i = nums.indexOf(k);
    const n = nums.length;
    const cnt = new Array((n << 1) | 1).fill(0);
    let ans = 1;
    let x = 0;
    for (let j = i + 1; j < n; ++j) {
        x += nums[j] > k ? 1 : -1;
        ans += x >= 0 && x <= 1 ? 1 : 0;
        ++cnt[x + n];
    }
    x = 0;
    for (let j = i - 1; ~j; --j) {
        x += nums[j] > k ? 1 : -1;
        ans += x >= 0 && x <= 1 ? 1 : 0;
        ans += cnt[-x + n] + cnt[-x + 1 + n];
    }
    return ans;
}

// 方法一：前缀和
export function countSubarraysTwo(nums: number[], k: number): number {
    const n = nums.length;
    let kIndex = -1;
    for (let i = 0; i < n; i++) {
        if (nums[i] === k) {
            kIndex = i;
            break;
        }
    }
    let ans = 0;
    const counts = new Map();
    counts.set(0, 1);
    let sum = 0;
    for (let i = 0; i < n; i++) {
        sum += sign(nums[i] - k);
        if (i < kIndex) {
            counts.set(sum, (counts.get(sum) || 0) + 1);
        } else {
            const prev0 = (counts.get(sum) || 0);
            const prev1 = (counts.get(sum - 1) || 0);
            ans += prev0 + prev1;
        }
    }
    return ans;
}

const sign = (num) => {
    if (num === 0) {
        return 0;
    }
    return num > 0 ? 1 : -1;
};
